
from modeller import *
from modeller.automodel import *
from modeller.scripts import complete_pdb

log.verbose()
env = environ()

preys = ['prey3','prey4-1','prey5','prey13','prey15'] #the names of the initial .ali files; has to be manually changed

def loop_refinement():
    for i in preys:
        
        env.io.atom_files_directory = './:../atom_files'

        class MyLoop(loopmodel):
           
            def select_loop_atoms(self):
                
                return selection(self.residue_range('17', '44')) #this is the location of our peptide
            def special_restraints(self, aln):
                rsr = self.restraints

                if preys.index(i) == 0:
                    rsr.add(secondary_structure.alpha(self.residue_range('24:','44:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('50:','54:')))
                    rsr.add(secondary_structure.strand(self.residue_range('9:','14:')))


                elif preys.index(i) == 1:
                    rsr.add(secondary_structure.alpha(self.residue_range('7:','7:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('57:','69:')))
                    rsr.add(secondary_structure.strand(self.residue_range('12:','12:')))
                    rsr.add(secondary_structure.strand(self.residue_range('35:','41:')))
                    rsr.add(secondary_structure.strand(self.residue_range('50:','53:')))
                    rsr.add(secondary_structure.strand(self.residue_range('76:','79:')))
                    rsr.add(secondary_structure.strand(self.residue_range('85:','88:')))
                    
                elif preys.index(i) == 2:
                    rsr.add(secondary_structure.alpha(self.residue_range('25:','31:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('36:','41:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('57:','70:')))
                    rsr.add(secondary_structure.strand(self.residue_range('9:','14:')))
                    rsr.add(secondary_structure.strand(self.residue_range('50:','54:')))
                    rsr.add(secondary_structure.strand(self.residue_range('76:','79:')))
                    rsr.add(secondary_structure.strand(self.residue_range('85:','89:')))
                elif preys.index(i) == 3:
                    rsr.add(secondary_structure.alpha(self.residue_range('28:','41:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('55:','68:')))
                    rsr.add(secondary_structure.strand(self.residue_range('8:','14:')))
                    rsr.add(secondary_structure.strand(self.residue_range('48:','51:')))
                    rsr.add(secondary_structure.strand(self.residue_range('74:','77:')))
                    rsr.add(secondary_structure.strand(self.residue_range('83:','87:')))
                elif preys.index(i) == 4:
                    rsr.add(secondary_structure.alpha(self.residue_range('26:','28:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('31:','31:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('33:','33:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('37:','37:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('39:','40:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('55:','68:')))
                    rsr.add(secondary_structure.strand(self.residue_range('8:','14:')))
                    rsr.add(secondary_structure.strand(self.residue_range('48:','52:')))
                    rsr.add(secondary_structure.strand(self.residue_range('74:','77:')))
                    rsr.add(secondary_structure.strand(self.residue_range('83:','87:')))

               

        m = MyLoop(env,
                   inimodel=str(i)+ ".pdb", # initial model of the target
                   sequence=str(i))          # code of the target

        m.loop.starting_model= 1           # index of the first loop model 
        m.loop.ending_model  = 10          # index of the last loop model
        m.loop.md_level = refine.slow # loop refinement method; this yields
                                           # models quickly but of low quality;
                                           # use refine.slow for better models

        m.make()

def evaluate_model():
    
    log.verbose()    
    env = environ()
    env.libs.topology.read(file='$(LIB)/top_heav.lib') 
    env.libs.parameters.read(file='$(LIB)/par.lib') 
    
    for x in preys:
        

        

        for i in range(1, 11):
           
            if i < 10:
                code = str(x)+(".BL000%s0001.pdb" % i)
            elif i >= 10:
                code = str(x)+(".BL00%s0001.pdb" % i)
            mdl = complete_pdb(env, code)
            s = selection(mdl)
            s.assess_dope(output='ENERGY_PROFILE NO_REPORT', file=code+'.profile',
                          normalize_profile=True, smoothing_window=15)

   
        
loop_refinement()
evaluate_model()
